package com.tianji.memory.mongo;

import cn.hutool.core.collection.CollStreamUtil;
import com.mongodb.client.result.UpdateResult;
import com.tianji.memory.OptimizeChatMemoryRepository;
import com.tianji.memory.pojo.ChatRecordDocument;
import com.tianji.memory.util.MessageUtil;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.data.redis.core.BoundListOperations;
import org.springframework.util.Assert;

import java.util.List;

/**
 * 基于 mongodb 的 ChatMemoryRepository
 */
public class MongoDBChatMemoryRepository implements ChatMemoryRepository, OptimizeChatMemoryRepository {

    @Resource
    private MongoTemplate mongoTemplate;

    private static final String COLLECTION = "chat_record";

    @Override
    public List<String> findConversationIds() {
        Query query = new Query();
        query.fields().include("conversationId");
        List<ChatRecordDocument> records = mongoTemplate.find(query, ChatRecordDocument.class);
        return CollStreamUtil.toList(records,  ChatRecordDocument::getConversationId);
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        ChatRecordDocument record = mongoTemplate.findOne(new Query(Criteria.where("conversationId").is(conversationId)), ChatRecordDocument.class);
        if (record == null) {
            return List.of();
        }
        return CollStreamUtil.toList(record.getMessages(), MessageUtil::toMessage);
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.notEmpty(messages, "消息列表不能为空");
        // 先删除
        deleteByConversationId(conversationId);
        // 再保存所有
        ChatRecordDocument record = ChatRecordDocument.builder()
                .conversationId(conversationId) // 设置对话id
                .messages(CollStreamUtil.toList(messages,MessageUtil::toJson)) // 设置当前对话下的所有对话数据
                .build();
        mongoTemplate.save(record);
    }



    @Override
    public void deleteByConversationId(String conversationId) {
        mongoTemplate.remove(new Query(Criteria.where("conversationId").is(conversationId)), ChatRecordDocument.class);
    }

    // TODO  mongoDB会话记录删除没记录好  需要整理
    @Override
    public void optimization(String conversationId) {
        Query query = Query.query(Criteria.where("conversationId").is(conversationId));
        ChatRecordDocument record = mongoTemplate.findOne(new Query(Criteria.where("conversationId").is(conversationId)), ChatRecordDocument.class);
        Update update = new Update();
        update.pop("messages", Update.Position.LAST); // 移除第1个元素
        update.pop("messages", Update.Position.LAST); // 再移除1个（原第2个元素，现在变成第1个）
        ChatRecordDocument record2 = mongoTemplate.findOne(new Query(Criteria.where("conversationId").is(conversationId)), ChatRecordDocument.class);
        UpdateResult updateResult = mongoTemplate.updateFirst(query, update, COLLECTION);
    }
}